from flask import Flask, request, jsonify, render_template, send_file
import redis
import os
import random

app = Flask(__name__)

# 连接到Redis服务器
r = redis.Redis(host='localhost', port=6379)

FILE_STORAGE_PATH = '/path/to/files/'

def generate_pickup_code():
    """生成一个6位数的取件码"""
    return ''.join([str(random.randint(0, 9)) for _ in range(6)])

def store_text_in_redis(text):
    """存储文本到Redis"""
    pickup_code = generate_pickup_code()
    text_hash = {
        'type': 'text',
        'content': text
    }
    r.hset(pickup_code, mapping=text_hash)
    r.expire(pickup_code, 3600)  # 1小时后自动删除
    return pickup_code

def store_file_in_filesystem(filename, content):
    """存储文件到文件系统并将元数据存储到Redis"""
    pickup_code = generate_pickup_code()
    filepath = os.path.join(FILE_STORAGE_PATH, pickup_code + "_" + filename)
    with open(filepath, 'wb') as f:
        f.write(content)
    
    file_hash = {
        'type': 'file',
        'filename': filename,
        'filepath': filepath
    }
    r.hset(pickup_code, mapping=file_hash)
    r.expire(pickup_code, 3600)  # 1小时后自动删除
    return pickup_code

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/store_text', methods=['POST'])
def store_text():
    content = request.json.get('content')
    pickup_code = store_text_in_redis(content)
    return jsonify({'pickup_code': pickup_code})

@app.route('/store_file', methods=['POST'])
def store_file():
    file = request.files['file']
    filename = file.filename
    content = file.read()
    pickup_code = store_file_in_filesystem(filename, content)
    return jsonify({'pickup_code': pickup_code})

@app.route('/pickup/<pickup_code>', methods=['GET'])
def pickup_item(pickup_code):
    item = r.hgetall(pickup_code)
    if not item:
        return jsonify({'error': '取件码无效或已过期'}), 404

    item_type = item[b'type'].decode()
    if item_type == 'text':
        content = item[b'content'].decode()
        return jsonify({'type': 'text', 'content': content})
    elif item_type == 'file':
        filepath = item[b'filepath'].decode()
        filename = item[b'filename'].decode()
        return jsonify({'type': 'file', 'filename': filename, 'url': f'/download/{pickup_code}'})

@app.route('/download/<pickup_code>', methods=['GET'])
def download_file(pickup_code):
    item = r.hgetall(pickup_code)
    if not item or item[b'type'].decode() != 'file':
        return jsonify({'error': '无效的取件码'}), 404

    filepath = item[b'filepath'].decode()
    filename = item[b'filename'].decode()
    return send_file(filepath, as_attachment=True, download_name=filename)

if __name__ == "__main__":
    if not os.path.exists(FILE_STORAGE_PATH):
        os.makedirs(FILE_STORAGE_PATH)
    app.run(host='0.0.0.0', port=80)
